setOldClass("nifti")
setClass(
    Class = "neurocon",
    representation = representation(
        address = "externalptr",
        coords = "data.frame",
        func = "character",
        ref = "nifti"
    )
)
setIs("neurocon", "big.matrix")

# Default is to allow parellel computing if setup
# options(allowParellel=FALSE) if want otherwise
options(allowParallel=TRUE)

# we can have a neurocon() method
# will get mat with ni.cor
# args: x, mask=NULL, ref=NULL
#   x can be:
#       - neurots (in which case no mask or ref is needed) => lead to correlation
#       - nifti filename (mask and ref optional) => lead to correlation
#       - big.matrix (coords and ref optional)

neurocon <- function(x, backingprefix, mask=NULL, ref=NULL, verbose=T, ...) {
    if (is.character(x))
        x <- neurots(x, mask, ref)
    
    bigmat <- cor.neurots(x, backingprefix, verbose=verbose, ...)
    
    if (class(bigmat) == "neurocon")
        return(bigmat)
    else
        return(new("neurocon", address=bigmat@address, coords=x@coords, func=x@fname, ref=x@ref))
}

cor.neurots <- function(x, backingprefix, blocksize=1000, overwrite=F, verbose=T) {
    if (class(x) != "neurots")
        stop("Input must be a neurots object")
    if (!is.character(backingprefix))
        stop("backingprefix must be a character/string")

    # Setup paths
    backingprefix <- rmext(abspath(backingprefix))
    backingpath <- dirname(backingprefix)
    backingfile <- sprintf("%s.bin", basename(backingprefix))
    descriptorfile <- sprintf("%s.desc", basename(backingprefix))
    neuroconfile <- sprintf("%s.neurocon", basename(backingprefix))
    
    # Check exists
    if (file.exists(file.path(backingpath, backingfile))) {
        warning(sprintf("Output file-backed correlation matrix '%s' already exists", backingfile))
        if (!overwrite) {
            message("Returning previous version")
            return(attach.neurocon(backingprefix))
        }
    }
    
    # Create empty matrix
    nVoxs <- ncol(x)
    nTimePts <- nrow(x)
    nBlocks <- ceiling(nVoxs/blocksize)
    cormat <- big.matrix(nVoxs, nVoxs, init=1, type="double", backingpath=backingpath, backingfile=backingfile, descriptorfile=descriptorfile)
    
    if (verbose) {
        pb <- create.progressbar(nBlocks)
        update.progressbar(pb, 0)
    }
    
    # Scale Time Series
    scalemat <- scale(x@mat)
    
    # Compute cross-product in chunks
    lapply(1:nBlocks, function(i) {
        blockstart <- (i-1)*blocksize + 1
        blockend <- min(nVoxs, blockstart + blocksize - 1)
        cols <- blockstart:blockend
        
        cormat[,cols] <- t(crossprod(scalemat[,cols], scalemat)/(nTimePts-1))

        if(verbose) update.progressbar(pb, i)
    })
    
    # Save additional neurocon info
    dput(list(coords=x@coords, func=x@fname, ref=x@ref$fname), file.path(backingpath, neuroconfile))
    
    if (verbose)
        end.progressbar(pb)
    
    return(cormat)
}

attach.neurocon <- function(backingprefix) {
    # Setup Paths
    backingprefix <- rmext(abspath(backingprefix))
    descriptorfile <- sprintf("%s/%s.desc", dirname(backingprefix), basename(backingprefix))
    neuroconfile <- sprintf("%s/%s.neurocon", dirname(backingprefix), basename(backingprefix))
    
    bigmat <- attach.big.matrix(descriptorfile)
    nclist <- dget(neuroconfile)
    
    return(new("neurocon", address=bigmat@address, coords=nclist$coords, func=nclist$func, ref=nifti.image.read(nclist$ref)))
}

batch.neurocon <- function(inputs, outputs, mask, ref=NULL, verbose=T, diagnostics=F, ...) {
    if (length(inputs) != length(outputs)) stop("Different number of inputs than outputs")
    nSubjs <- length(inputs)
    
    # Load mask
    ## load nifti filename
    if (is.character(mask))
        mask <- nifti.image.read(abspath(mask))
    ## setup nifti object
    if (class(mask) == "nifti")
        coords <- niGetCoords(mask[,,])
    else
        coords <- mask
    ## confirm that mask now a data frame
    if (!is.data.frame(coords))
        stop("Unrecognized type/object for mask: ", class(coords))
    
    # Load reference
    if (is.null(ref) && class(mask) == "nifti")
        ref <- mask
    else if (is.character(ref))
        ref <- nifti.image.read(abspath(ref))
    else if (class(ref) != "nifti")
        stop("Unrecognized type/object for ref: ", class(ref))
    
    if (verbose) cat("Computing correlations\n")
    
    # Loop through participants
    
    ## function that will be called
    corfun <- function(i) {
        if (verbose) cat(sprintf("Running subject %03i", i), "\n")
        
        # Check input
        if (!file.exists(inputs[i])) {
            warning(sprintf("Could not find functional file '%s', skipping subject", inputs[i]))
            return(NULL)
        }
        
        memusage <- gc(verbose=F, reset=T)
        
        # Run correlation
        systime <- system.time(cormat <- neurocon(inputs[i], outputs[i], mask=mask, ref=ref, verbose=verbose, ...))
        
        memusage <- gc(verbose=F)
        
        # Diagnostics
        if (diagnostics) {
            cat("---\n")
            cat("System Time:\n")
            print(systime)
            cat("Memory Usage:\n")
            print(memusage)
            cat("---\n")
        }
        
        cat("\n")
        
        return(cormat)
    }
    
    if (getOption("allowParallel") && getDoParRegistered())
        corlist <- foreach(i=1:nSubjs, .packages=c("bigmemory", "Rniftilib")) %dopar% corfun(i)
    else
        corlist <- lapply(1:nSubjs, corfun)
    
    return(corlist)
}

batch.attach.neurocon <- function(prefixes) {
    lapply(prefixes, function(p) attach.neurocon(p))
}

setMethod('print', signature(x='neurocon'),
    function(x) {
        cat("BIG MATRIX:\n")
        cat(sprintf("Connectivity matrix that is %i x %i", nrow(x), ncol(x)), "\n")
        cat("\n")
        
        cat("COORDS:\n")
        cat(sprintf("dim: %i x %i x %i", max(x@coords[,1]), max(x@coords[,2]), max(x@coords[,3])), "\n")
        cat(sprintf("nonzero voxels: %i", sum(x@coords$vals)), "\n")
        print(head(x@coords))
        cat("\n")
        
        cat("ORIG FUNC FILE:\n")
        cat(x@func, "\n")
        cat("\n")
        
        cat("REF FILE:\n")
        cat(sprintf("name: %s\nsize: %i x %i x %i", x@ref$fname, dim(x@ref)[1], dim(x@ref)[2], dim(x@ref)[3]), "\n")
        cat("\n")        
    }
)

setMethod('show', 'neurocon', function(object) print(object))